package com.LeeCode;

import java.util.*;

/**
 * 等值距离和
 */

public class Code2615 {
    public static void main(String[] args) {
        int[] nums = {1, 3, 1, 1, 2};
        System.out.println(Arrays.toString(new Code2615().distance(nums)));
    }

    public long[] distance(int[] nums) {
        int n = nums.length;
        long[] arr = new long[n];

        // 按值分组，存储每个值对应的所有索引
        Map<Integer, List<Integer>> map = new HashMap<>();
        for (int i = 0; i < n; i++) {
            map.computeIfAbsent(nums[i], k -> new ArrayList<>()).add(i);
        }

        // 对于每个相同值的索引列表，计算距离和
        for (List<Integer> indices : map.values()) {
            if (indices.size() <= 1) continue;

            // 计算前缀和
            int size = indices.size();
            long[] prefixSum = new long[size + 1];
            for (int i = 0; i < size; i++) {
                prefixSum[i + 1] = prefixSum[i] + indices.get(i);
            }

            // 对于每个位置，利用前缀和快速计算距离和
            for (int i = 0; i < size; i++) {
                int idx = indices.get(i);
                long leftSum = (long) i * idx - prefixSum[i];
                long rightSum = (prefixSum[size] - prefixSum[i + 1]) - (long) (size - i - 1) * idx;
                arr[idx] = leftSum + rightSum;
            }
        }

        return arr;
    }
}
